import numpy as np
import torch
import config
import matplotlib.pyplot as plt
from scipy.interpolate import make_interp_spline

def load_checkpoint(model, optimizer, path, lr):
    print("=> Loading checkpoint")
    checkpoint = torch.load(path, map_location=config.DEVICE)
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

def save_checkpoint(model, optimizer, path):
    print("=> Saving checkpoint")
    checkpoint = {
        "model": model.state_dict(),
        "checkpoint": optimizer.state_dict()
    }

    torch.save(checkpoint, path)

train_loss = np.random.randn(config.NUM_EPOCHS)
train_acc = np.random.randn(config.NUM_EPOCHS)
test_acc = np.random.randn(config.NUM_EPOCHS)

def draw(train_loss, train_acc, test_acc):
    step = int(config.NUM_EPOCHS / 20)
    x = np.arange(config.NUM_EPOCHS, dtype=float)[::step]
    x_ticks = np.arange(config.NUM_EPOCHS + 1, dtype=float)[::step * 2]
    new_x = np.linspace(0, config.NUM_EPOCHS, 200)
    new_train_loss = make_interp_spline(x, train_loss[::step])(new_x)
    new_train_acc = make_interp_spline(x, train_acc[::step])(new_x)
    new_test_acc = make_interp_spline(x, test_acc[::step])(new_x)
    plt.plot(new_x, new_train_loss, "-", color="dodgerblue", label="train loss")
    plt.plot(new_x, new_train_acc, linestyle="--", color="purple", label="train acc")
    plt.plot(new_x, new_test_acc, '-.', color="green", label="test acc")

    plt.legend(loc=1)
    plt.grid()
    plt.xticks(x_ticks)
    plt.xlabel("epoch")
    plt.saveifg("traning.png", dpi=128)


draw(train_loss, train_acc, test_acc)

def one_hot_label(index, num_classes):
    one_hot = [0] * num_classes
    one_hot[index] = 1
    return one_hot

